#pragma once
#include "kernel_operator.h"

namespace AscendC {

constexpr int zhishubiao[500] = {
    2,    3,    5,    7,    11,   13,   17,   19,   23,   29,   31,   37,
    41,   43,   47,   53,   59,   61,   67,   71,   73,   79,   83,   89,
    97,   101,  103,  107,  109,  113,  127,  131,  137,  139,  149,  151,
    157,  163,  167,  173,  179,  181,  191,  193,  197,  199,  211,  223,
    227,  229,  233,  239,  241,  251,  257,  263,  269,  271,  277,  281,
    283,  293,  307,  311,  313,  317,  331,  337,  347,  349,  353,  359,
    367,  373,  379,  383,  389,  397,  401,  409,  419,  421,  431,  433,
    439,  443,  449,  457,  461,  463,  467,  479,  487,  491,  499,  503,
    509,  521,  523,  541,  547,  557,  563,  569,  571,  577,  587,  593,
    599,  601,  607,  613,  617,  619,  631,  641,  643,  647,  653,  659,
    661,  673,  677,  683,  691,  701,  709,  719,  727,  733,  739,  743,
    751,  757,  761,  769,  773,  787,  797,  809,  811,  821,  823,  827,
    829,  839,  853,  857,  859,  863,  877,  881,  883,  887,  907,  911,
    919,  929,  937,  941,  947,  953,  967,  971,  977,  983,  991,  997,
    1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069,
    1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163,
    1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249,
    1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321,
    1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439,
    1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511,
    1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601,
    1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657, 1663, 1667, 1669, 1693,
    1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783,
    1787, 1789, 1801, 1811, 1823, 1831, 1847, 1861, 1867, 1871, 1873, 1877,
    1879, 1889, 1901, 1907, 1913, 1931, 1933, 1949, 1951, 1973, 1979, 1987,
    1993, 1997, 1999, 2003, 2011, 2017, 2027, 2029, 2039, 2053, 2063, 2069,
    2081, 2083, 2087, 2089, 2099, 2111, 2113, 2129, 2131, 2137, 2141, 2143,
    2153, 2161, 2179, 2203, 2207, 2213, 2221, 2237, 2239, 2243, 2251, 2267,
    2269, 2273, 2281, 2287, 2293, 2297, 2309, 2311, 2333, 2339, 2341, 2347,
    2351, 2357, 2371, 2377, 2381, 2383, 2389, 2393, 2399, 2411, 2417, 2423,
    2437, 2441, 2447, 2459, 2467, 2473, 2477, 2503, 2521, 2531, 2539, 2543,
    2549, 2551, 2557, 2579, 2591, 2593, 2609, 2617, 2621, 2633, 2647, 2657,
    2659, 2663, 2671, 2677, 2683, 2687, 2689, 2693, 2699, 2707, 2711, 2713,
    2719, 2729, 2731, 2741, 2749, 2753, 2767, 2777, 2789, 2791, 2797, 2801,
    2803, 2819, 2833, 2837, 2843, 2851, 2857, 2861, 2879, 2887, 2897, 2903,
    2909, 2917, 2927, 2939, 2953, 2957, 2963, 2969, 2971, 2999, 3001, 3011,
    3019, 3023, 3037, 3041, 3049, 3061, 3067, 3079, 3083, 3089, 3109, 3119,
    3121, 3137, 3163, 3167, 3169, 3181, 3187, 3191, 3203, 3209, 3217, 3221,
    3229, 3251, 3253, 3257, 3259, 3271, 3299, 3301, 3307, 3313, 3319, 3323,
    3329, 3331, 3343, 3347, 3359, 3361, 3371, 3373, 3389, 3391, 3407, 3413,
    3433, 3449, 3457, 3461, 3463, 3467, 3469, 3491, 3499, 3511, 3517, 3527,
    3529, 3533, 3539, 3541, 3547, 3557, 3559, 3571};
constexpr int map_lengths[500] = {
    2,      5,      10,     17,     28,     41,     58,     77,     100,
    129,    160,    197,    238,    281,    328,    381,    440,    501,
    568,    639,    712,    791,    874,    963,    1060,   1161,   1264,
    1371,   1480,   1593,   1720,   1851,   1988,   2127,   2276,   2427,
    2584,   2747,   2914,   3087,   3266,   3447,   3638,   3831,   4028,
    4227,   4438,   4661,   4888,   5117,   5350,   5589,   5830,   6081,
    6338,   6601,   6870,   7141,   7418,   7699,   7982,   8275,   8582,
    8893,   9206,   9523,   9854,   10191,  10538,  10887,  11240,  11599,
    11966,  12339,  12718,  13101,  13490,  13887,  14288,  14697,  15116,
    15537,  15968,  16401,  16840,  17283,  17732,  18189,  18650,  19113,
    19580,  20059,  20546,  21037,  21536,  22039,  22548,  23069,  23592,
    24133,  24680,  25237,  25800,  26369,  26940,  27517,  28104,  28697,
    29296,  29897,  30504,  31117,  31734,  32353,  32984,  33625,  34268,
    34915,  35568,  36227,  36888,  37561,  38238,  38921,  39612,  40313,
    41022,  41741,  42468,  43201,  43940,  44683,  45434,  46191,  46952,
    47721,  48494,  49281,  50078,  50887,  51698,  52519,  53342,  54169,
    54998,  55837,  56690,  57547,  58406,  59269,  60146,  61027,  61910,
    62797,  63704,  64615,  65534,  66463,  67400,  68341,  69288,  70241,
    71208,  72179,  73156,  74139,  75130,  76127,  77136,  78149,  79168,
    80189,  81220,  82253,  83292,  84341,  85392,  86453,  87516,  88585,
    89672,  90763,  91856,  92953,  94056,  95165,  96282,  97405,  98534,
    99685,  100838, 102001, 103172, 104353, 105540, 106733, 107934, 109147,
    110364, 111587, 112816, 114047, 115284, 116533, 117792, 119069, 120348,
    121631, 122920, 124211, 125508, 126809, 128112, 129419, 130738, 132059,
    133386, 134747, 136114, 137487, 138868, 140267, 141676, 143099, 144526,
    145955, 147388, 148827, 150274, 151725, 153178, 154637, 156108, 157589,
    159072, 160559, 162048, 163541, 165040, 166551, 168074, 169605, 171148,
    172697, 174250, 175809, 177376, 178947, 180526, 182109, 183706, 185307,
    186914, 188523, 190136, 191755, 193376, 195003, 196640, 198297, 199960,
    201627, 203296, 204989, 206686, 208385, 210094, 211815, 213538, 215271,
    217012, 218759, 220512, 222271, 224048, 225831, 227618, 229407, 231208,
    233019, 234842, 236673, 238520, 240381, 242248, 244119, 245992, 247869,
    249748, 251637, 253538, 255445, 257358, 259289, 261222, 263171, 265122,
    267095, 269074, 271061, 273054, 275051, 277050, 279053, 281064, 283081,
    285108, 287137, 289176, 291229, 293292, 295361, 297442, 299525, 301612,
    303701, 305800, 307911, 310024, 312153, 314284, 316421, 318562, 320705,
    322858, 325019, 327198, 329401, 331608, 333821, 336042, 338279, 340518,
    342761, 345012, 347279, 349548, 351821, 354102, 356389, 358682, 360979,
    363288, 365599, 367932, 370271, 372612, 374959, 377310, 379667, 382038,
    384415, 386796, 389179, 391568, 393961, 396360, 398771, 401188, 403611,
    406048, 408489, 410936, 413395, 415862, 418335, 420812, 423315, 425836,
    428367, 430906, 433449, 435998, 438549, 441106, 443685, 446276, 448869,
    451478, 454095, 456716, 459349, 461996, 464653, 467312, 469975, 472646,
    475323, 478006, 480693, 483382, 486075, 488774, 491481, 494192, 496905,
    499624, 502353, 505084, 507825, 510574, 513327, 516094, 518871, 521660,
    524451, 527248, 530049, 532852, 535671, 538504, 541341, 544184, 547035,
    549892, 552753, 555632, 558519, 561416, 564319, 567228, 570145, 573072,
    576011, 578964, 581921, 584884, 587853, 590824, 593823, 596824, 599835,
    602854, 605877, 608914, 611955, 615004, 618065, 621132, 624211, 627294,
    630383, 633492, 636611, 639732, 642869, 646032, 649199, 652368, 655549,
    658736, 661927, 665130, 668339, 671556, 674777, 678006, 681257, 684510,
    687767, 691026, 694297, 697596, 700897, 704204, 707517, 710836, 714159,
    717488, 720819, 724162, 727509, 730868, 734229, 737600, 740973, 744362,
    747753, 751160, 754573, 758006, 761455, 764912, 768373, 771836, 775303,
    778772, 782263, 785762, 789273, 792790, 796317, 799846, 803379, 806918,
    810459, 814006, 817563, 821122, 824693,
};
constexpr int map_start_index[500] = {
    0,      2,      5,      10,     17,     28,     41,     58,     77,
    100,    129,    160,    197,    238,    281,    328,    381,    440,
    501,    568,    639,    712,    791,    874,    963,    1060,   1161,
    1264,   1371,   1480,   1593,   1720,   1851,   1988,   2127,   2276,
    2427,   2584,   2747,   2914,   3087,   3266,   3447,   3638,   3831,
    4028,   4227,   4438,   4661,   4888,   5117,   5350,   5589,   5830,
    6081,   6338,   6601,   6870,   7141,   7418,   7699,   7982,   8275,
    8582,   8893,   9206,   9523,   9854,   10191,  10538,  10887,  11240,
    11599,  11966,  12339,  12718,  13101,  13490,  13887,  14288,  14697,
    15116,  15537,  15968,  16401,  16840,  17283,  17732,  18189,  18650,
    19113,  19580,  20059,  20546,  21037,  21536,  22039,  22548,  23069,
    23592,  24133,  24680,  25237,  25800,  26369,  26940,  27517,  28104,
    28697,  29296,  29897,  30504,  31117,  31734,  32353,  32984,  33625,
    34268,  34915,  35568,  36227,  36888,  37561,  38238,  38921,  39612,
    40313,  41022,  41741,  42468,  43201,  43940,  44683,  45434,  46191,
    46952,  47721,  48494,  49281,  50078,  50887,  51698,  52519,  53342,
    54169,  54998,  55837,  56690,  57547,  58406,  59269,  60146,  61027,
    61910,  62797,  63704,  64615,  65534,  66463,  67400,  68341,  69288,
    70241,  71208,  72179,  73156,  74139,  75130,  76127,  77136,  78149,
    79168,  80189,  81220,  82253,  83292,  84341,  85392,  86453,  87516,
    88585,  89672,  90763,  91856,  92953,  94056,  95165,  96282,  97405,
    98534,  99685,  100838, 102001, 103172, 104353, 105540, 106733, 107934,
    109147, 110364, 111587, 112816, 114047, 115284, 116533, 117792, 119069,
    120348, 121631, 122920, 124211, 125508, 126809, 128112, 129419, 130738,
    132059, 133386, 134747, 136114, 137487, 138868, 140267, 141676, 143099,
    144526, 145955, 147388, 148827, 150274, 151725, 153178, 154637, 156108,
    157589, 159072, 160559, 162048, 163541, 165040, 166551, 168074, 169605,
    171148, 172697, 174250, 175809, 177376, 178947, 180526, 182109, 183706,
    185307, 186914, 188523, 190136, 191755, 193376, 195003, 196640, 198297,
    199960, 201627, 203296, 204989, 206686, 208385, 210094, 211815, 213538,
    215271, 217012, 218759, 220512, 222271, 224048, 225831, 227618, 229407,
    231208, 233019, 234842, 236673, 238520, 240381, 242248, 244119, 245992,
    247869, 249748, 251637, 253538, 255445, 257358, 259289, 261222, 263171,
    265122, 267095, 269074, 271061, 273054, 275051, 277050, 279053, 281064,
    283081, 285108, 287137, 289176, 291229, 293292, 295361, 297442, 299525,
    301612, 303701, 305800, 307911, 310024, 312153, 314284, 316421, 318562,
    320705, 322858, 325019, 327198, 329401, 331608, 333821, 336042, 338279,
    340518, 342761, 345012, 347279, 349548, 351821, 354102, 356389, 358682,
    360979, 363288, 365599, 367932, 370271, 372612, 374959, 377310, 379667,
    382038, 384415, 386796, 389179, 391568, 393961, 396360, 398771, 401188,
    403611, 406048, 408489, 410936, 413395, 415862, 418335, 420812, 423315,
    425836, 428367, 430906, 433449, 435998, 438549, 441106, 443685, 446276,
    448869, 451478, 454095, 456716, 459349, 461996, 464653, 467312, 469975,
    472646, 475323, 478006, 480693, 483382, 486075, 488774, 491481, 494192,
    496905, 499624, 502353, 505084, 507825, 510574, 513327, 516094, 518871,
    521660, 524451, 527248, 530049, 532852, 535671, 538504, 541341, 544184,
    547035, 549892, 552753, 555632, 558519, 561416, 564319, 567228, 570145,
    573072, 576011, 578964, 581921, 584884, 587853, 590824, 593823, 596824,
    599835, 602854, 605877, 608914, 611955, 615004, 618065, 621132, 624211,
    627294, 630383, 633492, 636611, 639732, 642869, 646032, 649199, 652368,
    655549, 658736, 661927, 665130, 668339, 671556, 674777, 678006, 681257,
    684510, 687767, 691026, 694297, 697596, 700897, 704204, 707517, 710836,
    714159, 717488, 720819, 724162, 727509, 730868, 734229, 737600, 740973,
    744362, 747753, 751160, 754573, 758006, 761455, 764912, 768373, 771836,
    775303, 778772, 782263, 785762, 789273, 792790, 796317, 799846, 803379,
    806918, 810459, 814006, 817563, 821122,
};

template <typename KT, typename VT> class IntHashNode {
public:
  KT key = 0;
  VT value = 0;
  bool hasValue = false;
};

template <int length, typename KT, typename VT> class IntHashMap {
public:
  IntHashNode<KT, VT> data_[length];
  int64_t size_ = 0;
  __aicore__ inline IntHashMap() {}
  __aicore__ inline int key2Index(KT key, int zhishuIndex) {
    return map_start_index[zhishuIndex] + key % zhishubiao[zhishuIndex];
  }
  __aicore__ inline VT getNodeIndex(KT key) {
    for (int i = 0; i < 500; ++i) {
      int index = key2Index(key, i);
      if (index < length) {
        auto node = data_ + index;
        if (node->hasValue && node->key == key) {
          return index;
        }
      } else {
        break;
      }
    }

    for (int i = length - 1; i >= 0; --i) {
      auto node = data_ + i;
      if (node->hasValue && node->key == key) {
        return i;
      }
    }

    return -1;
  }

  __aicore__ inline VT get(KT key, VT default_value = 0) {
    auto index = getNodeIndex(key);
    if (index == -1) {
      return default_value;
    }
    return data_[index].value;
  }

  __aicore__ inline void remove(KT key) {
    auto index = getNodeIndex(key);
    if (index == -1) {
      return;
    }
    data_[index].hasValue = false;
    size_--;
  }

  __aicore__ inline bool trySetNode(IntHashNode<KT, VT> *node, KT key,
                                    VT value) {
    if (!node->hasValue) {
      node->key = key;
      node->value = value;
      node->hasValue = true;
      size_++;
      return true;
    }
    if (node->key == key) {
      node->value = value;
      return true;
    }
    return false;
  }

  __aicore__ inline bool set(KT key, VT value) {
    for (int i = 0; i < 500; ++i) {
      int index = key2Index(key, i);
      if (index < length) {
        if (trySetNode(data_ + index, key, value)) {
          return true;
        }
      } else {
        break;
      }
    }

    for (int i = length - 1; i >= 0; --i) {
      if (trySetNode(data_ + i, key, value)) {
        return true;
      }
    }

    return false;
  }

  template <typename T> __aicore__ inline void print_(T name) {
    printf("--------------------%s---------------------\n", name);
    for (int i = 0; i < length; ++i) {
      auto node = data_[i];
      if (node.hasValue) {
        printf("%d : (%d, %d) \n", i, node.key, node.value);
      }
    }
    printf("--------------------%s end-----------------\n", name);
  }
};

} // namespace AscendC
